import logging
from typing import Union, List, Dict, Any
import pymysql
from pymysql.cursors import DictCursor
from config import config


class MysqlDB:
    """MySQL数据库操作封装类"""
    
    def __init__(self) -> None:
        try:
            # 连接数据库
            self.connection = pymysql.connect(
                host=config.MysqlDB.host,
                user=config.MysqlDB.username,  # 修正参数名
                password=config.MysqlDB.password,
                port=config.MysqlDB.port,
                database=config.MysqlDB.database,  # 添加数据库名
                charset='utf8mb4',
                cursorclass=DictCursor  # 使用字典游标
            )
            # 创建游标
            self.cursor = self.connection.cursor()
            
        except Exception as e:
            logging.error("数据库连接失败，失败原因: %s", e)
            raise

    def close(self) -> None:
        """关闭数据库连接"""
        try:
            self.cursor.close()
            self.connection.close()
        except Exception as e:
            logging.error("关闭数据库连接失败: %s", e)
            raise

    def query(self, sql: str, params: tuple = None) -> List[Dict]:
        """
        执行查询操作
        :param sql: SQL查询语句
        :param params: 查询参数
        :return: 查询结果列表
        """
        try:
            self.cursor.execute(sql, params)
            return self.cursor.fetchall()
        except Exception as e:
            logging.error("查询失败，SQL: %s, 参数: %s, 错误: %s", sql, params, e)
            raise

    def execute(self, sql: str, params: tuple = None) -> int:
        """
        执行增删改操作
        :param sql: SQL语句
        :param params: 执行参数
        :return: 影响行数
        """
        try:
            rows = self.cursor.execute(sql, params)
            self.connection.commit()
            return rows
        except Exception as e:
            self.connection.rollback()
            logging.error("执行SQL失败，SQL: %s, 参数: %s, 错误: %s", sql, params, e)
            raise

    def insert(self, table: str, data: Dict[str, Any]) -> int:
        """
        插入数据
        :param table: 表名
        :param data: 插入的数据字典
        :return: 插入成功的行数
        """
        fields = ','.join(data.keys())
        values = ','.join(['%s'] * len(data))
        sql = f"INSERT INTO {table} ({fields}) VALUES ({values})"
        return self.execute(sql, tuple(data.values()))

    def update(self, table: str, data: Dict[str, Any], condition: str, params: tuple = None) -> int:
        """
        更新数据
        :param table: 表名
        :param data: 更新的数据字典
        :param condition: WHERE条件
        :param params: 条件参数
        :return: 更新成功的行数
        """
        set_items = ','.join([f"{k}=%s" for k in data.keys()])
        sql = f"UPDATE {table} SET {set_items} WHERE {condition}"
        all_params = tuple(data.values()) + (params or ())
        return self.execute(sql, all_params)

    def delete(self, table: str, condition: str, params: tuple = None) -> int:
        """
        删除数据
        :param table: 表名
        :param condition: WHERE条件
        :param params: 条件参数
        :return: 删除成功的行数
        """
        sql = f"DELETE FROM {table} WHERE {condition}"
        return self.execute(sql, params)

    def assert_data(self, table: str, expected_data: Dict[str, Any], condition: str, 
                    params: tuple = None) -> bool:
        """
        断言数据是否符合预期
        :param table: 表名
        :param expected_data: 预期的数据字典
        :param condition: WHERE条件
        :param params: 条件参数
        :return: 断言结果
        """
        fields = ','.join(expected_data.keys())
        sql = f"SELECT {fields} FROM {table} WHERE {condition}"
        actual_data = self.query(sql, params)
        
        if not actual_data:
            logging.error("未找到符合条件的数据")
            return False
            
        # 比较第一条记录与预期数据
        actual_record = actual_data[0]
        for key, expected_value in expected_data.items():
            if key not in actual_record:
                logging.error(f"字段 {key} 不存在于查询结果中")
                return False
            if actual_record[key] != expected_value:
                logging.error(f"字段 {key} 的值不匹配: 预期 {expected_value}, 实际 {actual_record[key]}")
                return False
                
        return True

    def __enter__(self):
        """支持with语句"""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """退出时自动关闭连接"""
        self.close()